import pandas as pd
from L2.memory_module import MemoryModule
import numpy as np
import random

# 测试代码
if __name__ == '__main__':
    # 假设的环境状态和动作维度
    state_dim = 1
    action_dim = 1

    # 创建回放缓冲区
    buffer = MemoryModule(capacity=250, file_name='memory.csv', alpha=0.6)

    # 模拟生成和添加经验
    for _ in range(100):
        state_test = np.random.randn(state_dim)  # 随机生成当前状态
        action_test = np.random.randint(action_dim)  # 随机生成一个动作
        reward_test = random.random()  # 随机生成一个奖励
        next_state_test = np.random.randn(state_dim)  # 随机生成下一个状态
        done_test = np.random.choice([True, False])  # 随机生成结束标志
        id_test = np.random.randint(0, 100)
        importance_test = np.random.randint(0, 10)
        type_test = np.random.choice(['work', 'study', 'play', 'communication', 'family'])
        buffer.push(
            id_test, importance_test, type_test,
            [state_test, action_test, reward_test, next_state_test, done_test]
        )  # 将经验添加到缓冲区
        # print(id_test)
        # buffer.push(1, "test", 22)  # 将经验添加到缓冲区

    buffer.save_memory()

    # 从缓冲区中采样
    samples_test = buffer.priority_weighted_sample(batch_size=32)  # 使用均匀随机采样策略
    buffer.forget_by_time(10)

    # 假设从某个学习过程中获得了新的优先级
    # new_priorities = np.random.random(size=32)

    # 假设的采样索引（在实际应用中，这些索引应从sample方法返回）
    # sample_indices = np.random.choice(range(len(buffer)), 32, replace=False)

    # 更新采样经验的优先级
    # buffer.update_priorities(sample_indices, new_priorities)




